package org.apache.flink.connector.scylladb.table;

import org.apache.flink.connector.scylladb.source.ScyllaDBLookupFunction;
import org.apache.flink.connector.scylladb.source.ScyllaDBSourceFunction;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.connector.ChangelogMode;
import org.apache.flink.table.connector.source.*;
import org.apache.flink.table.connector.source.abilities.SupportsLimitPushDown;
import org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.RowType;

import java.time.Duration;
import java.util.Arrays;

/**
 * ScyllaDB 动态表源实现，支持扫描和查找操作
 */
public class ScyllaDBDynamicTableSource implements ScanTableSource, LookupTableSource, 
        SupportsProjectionPushDown, SupportsLimitPushDown {

    private final String url;
    private final String keySpace;
    private final String tableName;
    private final String username;
    private final String password;
    private final int scanFetchSize;
    private final String scanPartitionColumn;
    private final int scanParallelism;
    private final long lookupCacheMaxRows;
    private final Duration lookupCacheTtl;
    private final int lookupMaxRetries;
    private final int connectionPoolLocalSize;
    private final int connectionPoolRemoteSize;
    private final Duration connectionTimeout;
    private final Duration requestTimeout;
    private final String consistencyLevel;
    private final boolean sslEnabled;
    private final String sslKeystorePath;
    private final String sslKeystorePassword;
    private final String sslTruststorePath;
    private final String sslTruststorePassword;
    private final boolean lookupAggregateMode;
    private final String lookupAggregateFieldName;
    private final DataType physicalDataType;
    private final ResolvedSchema resolvedSchema;

    // 投影相关字段
    private int[][] projectedFields;
    private long limit = -1;

    public ScyllaDBDynamicTableSource(
            String url,
            String keySpace,
            String tableName,
            String username,
            String password,
            int scanFetchSize,
            String scanPartitionColumn,
            int scanParallelism,
            long lookupCacheMaxRows,
            Duration lookupCacheTtl,
            int lookupMaxRetries,
            int connectionPoolLocalSize,
            int connectionPoolRemoteSize,
            Duration connectionTimeout,
            Duration requestTimeout,
            String consistencyLevel,
            boolean sslEnabled,
            String sslKeystorePath,
            String sslKeystorePassword,
            String sslTruststorePath,
            String sslTruststorePassword,
            boolean lookupAggregateMode,
            String lookupAggregateFieldName,
            DataType physicalDataType,
            ResolvedSchema resolvedSchema) {
        
        this.url = url;
        this.keySpace = keySpace;
        this.tableName = tableName;
        this.username = username;
        this.password = password;
        this.scanFetchSize = scanFetchSize;
        this.scanPartitionColumn = scanPartitionColumn;
        this.scanParallelism = scanParallelism;
        this.lookupCacheMaxRows = lookupCacheMaxRows;
        this.lookupCacheTtl = lookupCacheTtl;
        this.lookupMaxRetries = lookupMaxRetries;
        this.connectionPoolLocalSize = connectionPoolLocalSize;
        this.connectionPoolRemoteSize = connectionPoolRemoteSize;
        this.connectionTimeout = connectionTimeout;
        this.requestTimeout = requestTimeout;
        this.consistencyLevel = consistencyLevel;
        this.sslEnabled = sslEnabled;
        this.sslKeystorePath = sslKeystorePath;
        this.sslKeystorePassword = sslKeystorePassword;
        this.sslTruststorePath = sslTruststorePath;
        this.sslTruststorePassword = sslTruststorePassword;
        this.lookupAggregateMode = lookupAggregateMode;
        this.lookupAggregateFieldName = lookupAggregateFieldName;
        this.physicalDataType = physicalDataType;
        this.resolvedSchema = resolvedSchema;
    }

    @Override
    public ChangelogMode getChangelogMode() {
        // ScyllaDB 支持 INSERT 操作，主要用于批量读取
        return ChangelogMode.insertOnly();
    }

    @Override
    public ScanRuntimeProvider getScanRuntimeProvider(ScanContext context) {
        final RowType physicalRowType = (RowType) physicalDataType.getLogicalType();
        
        final ScyllaDBSourceFunction sourceFunction = new ScyllaDBSourceFunction(
                url,
                keySpace,
                tableName,
                username,
                password,
                scanFetchSize,
                scanPartitionColumn,
                connectionPoolLocalSize,
                connectionPoolRemoteSize,
                connectionTimeout,
                requestTimeout,
                consistencyLevel,
                sslEnabled,
                sslKeystorePath,
                sslKeystorePassword,
                sslTruststorePath,
                sslTruststorePassword,
                physicalRowType,
                projectedFields,
                limit
        );

        return SourceFunctionProvider.of(sourceFunction, false);
    }

    @Override
    public LookupRuntimeProvider getLookupRuntimeProvider(LookupContext context) {
        final RowType physicalRowType = (RowType) physicalDataType.getLogicalType();
        
        // 处理复合键：提取所有键字段的索引
        int[][] lookupKeys = context.getKeys();
        int[] keyIndices = new int[lookupKeys.length];
        for (int i = 0; i < lookupKeys.length; i++) {
            keyIndices[i] = lookupKeys[i][0];
        }
        
        // 如果是聚合模式，创建特殊的 LookupFunction
        if (lookupAggregateMode) {
            final ScyllaDBLookupFunction lookupFunction = new ScyllaDBLookupFunction(
                    url,
                    keySpace,
                    tableName,
                    username,
                    password,
                    lookupCacheMaxRows,
                    lookupCacheTtl,
                    lookupMaxRetries,
                    connectionPoolLocalSize,
                    connectionPoolRemoteSize,
                    connectionTimeout,
                    requestTimeout,
                    consistencyLevel,
                    sslEnabled,
                    sslKeystorePath,
                    sslKeystorePassword,
                    sslTruststorePath,
                    sslTruststorePassword,
                    physicalRowType,
                    keyIndices,
                    lookupAggregateMode,
                    lookupAggregateFieldName
            );
            return TableFunctionProvider.of(lookupFunction);
        } else {
            // 普通模式
            final ScyllaDBLookupFunction lookupFunction = new ScyllaDBLookupFunction(
                    url,
                    keySpace,
                    tableName,
                    username,
                    password,
                    lookupCacheMaxRows,
                    lookupCacheTtl,
                    lookupMaxRetries,
                    connectionPoolLocalSize,
                    connectionPoolRemoteSize,
                    connectionTimeout,
                    requestTimeout,
                    consistencyLevel,
                    sslEnabled,
                    sslKeystorePath,
                    sslKeystorePassword,
                    sslTruststorePath,
                    sslTruststorePassword,
                    physicalRowType,
                    keyIndices,
                    false, // 非聚合模式
                    null   // 不需要聚合字段名
            );
            return TableFunctionProvider.of(lookupFunction);
        }
    }

    @Override
    public boolean supportsNestedProjection() {
        return false;
    }

    @Override
    public void applyProjection(int[][] projectedFields, DataType producedDataType) {
        this.projectedFields = projectedFields;
    }

    @Override
    public void applyLimit(long limit) {
        this.limit = limit;
    }

    @Override
    public DynamicTableSource copy() {
        final ScyllaDBDynamicTableSource copy = new ScyllaDBDynamicTableSource(
                url,
                keySpace,
                tableName,
                username,
                password,
                scanFetchSize,
                scanPartitionColumn,
                scanParallelism,
                lookupCacheMaxRows,
                lookupCacheTtl,
                lookupMaxRetries,
                connectionPoolLocalSize,
                connectionPoolRemoteSize,
                connectionTimeout,
                requestTimeout,
                consistencyLevel,
                sslEnabled,
                sslKeystorePath,
                sslKeystorePassword,
                sslTruststorePath,
                sslTruststorePassword,
                lookupAggregateMode,
                lookupAggregateFieldName,
                physicalDataType,
                resolvedSchema
        );
        copy.projectedFields = this.projectedFields;
        copy.limit = this.limit;
        return copy;
    }

    @Override
    public String asSummaryString() {
        return String.format("ScyllaDB(url=%s, keyspace=%s, table=%s)", 
                url, keySpace, tableName);
    }
} 